from python_ai.category.data.mnist_dataset import x_train, x_test, x_val, y_train, y_test, y_val, shape_, y_train_oh, y_test_oh, y_val_oh
from python_ai.common.xcommon import *
import matplotlib.pyplot as plt
import os
import io
from PIL import Image
from mlp_impl_separated_bias_from_weights import MLPsb
import mlp_impl_sigmoid as mlps
import python_ai.category.tvts.tvts as tvts
import sys
from python_ai.category.NumpyNet.utils import *

sep('Prepare data')
x_train = uint8_to_flt_by_lut(x_train)
x_val = uint8_to_flt_by_lut(x_val)
x_test = uint8_to_flt_by_lut(x_test)
dims = shape_[0] * shape_[1]
x_train = x_train.reshape(-1, dims)
x_val = x_val.reshape(-1, dims)
x_test = x_test.reshape(-1, dims)
y_train = y_train.reshape(-1, 1)
y_val = y_val.reshape(-1, 1)
y_test = y_test.reshape(-1, 1)
print('x_train', x_train.shape)
print('x_val', x_val.shape)
print('x_test', x_test.shape)
print('y_train', y_train.shape)
print('y_val', y_val.shape)
print('y_test', y_test.shape)
print('y_train_oh', y_train_oh.shape)
print('y_val_oh', y_val_oh.shape)
print('y_test_oh', y_test_oh.shape)
print('x_train', x_train[0, :100])
print('x_val', x_val[0, :100])
print('x_test', x_test[0, :100])
print('y_train', np.unique(y_train))
print('y_val', np.unique(y_val))
print('y_test', np.unique(y_test))
# sys.exit(0)

# Hyper params
NAME = 'mnist_multi_clf_by_numpynet_mlp_softmax_ok_sep_bias'
MEMO = 'Full training data set, uint8 to flt'
RSEED = None
if RSEED is not None:
    np.random.seed(RSEED)
VER = 'v1.0'
N_BATCH_SIZE = 256
LR = 0.1  # learning rate (It should be higher for all sigmoid layers than all relu ones)
N_EPOCH = 200  # Please to be greater than 650 for a good model; to be 300 for a normal P-R curve
SAVE_FREQ = 20
BASE_DIR, FILE_NAME = os.path.split(__file__)
SAVE_DIR = os.path.join(BASE_DIR, '_save', FILE_NAME, VER)
os.makedirs(SAVE_DIR, exist_ok=True)

#ts = tvts.tvts(NAME, '172.26.41.157', init_params={
ts = tvts.tvts(NAME, params={
    'ver': VER,
    'batch_size': N_BATCH_SIZE,
    'lr': LR,
    'n_epoch': N_EPOCH,
    'memo': MEMO
}, save_freq=SAVE_FREQ)
CKPT_PATH = ts.resume(7, 0)
IS_TRAIN = 0


def get_saved_path_prefix(xepoch):
    return os.path.join(SAVE_DIR, ts.get_save_prefix(xepoch))


def get_saved_path(xepoch):
    prefix = get_saved_path_prefix(xepoch)
    return f'{prefix}.npy'


def get_curve_path(xsave_path):
    return f'{xsave_path}.png'


def run_on_data(model, label, x_data, y_data, n_batch_size, is_train, epoch):
    n = len(x_data)
    n_batch = int(np.ceil(n / n_batch_size))
    cost_avg = 0.
    metrics_avg = None
    pred = np.array([])
    for i, (x_batch, y_batch) in enumerate(mlps.dl(x_data, y_data, n_batch_size, is_train)):
        batch = i + 1
        h = model(x_batch)
        pred = np.append(pred, h)
        cost = model.cost(h, y_batch)
        metrics = model.metrics(y_batch, h)
        print(label, 'epoch:', epoch, 'batch:', batch, 'cost:', cost, 'metrics:', metrics)
        acc = metrics.tolist()[0][0]
        if is_train:
            ts.save_batch(epoch, batch, {
                'cost': cost,
                'acc': acc
            })
        cost_avg += cost / n_batch
        if metrics_avg is None:
            metrics_avg = np.zeros_like(metrics, dtype=np.float32)
        metrics_avg += metrics / n_batch
        if is_train:
            model.backward(y_batch)
            model.step()
    print(label, 'epoch:', epoch, 'cost_avg:', cost_avg, 'metrics:', metrics_avg)
    return cost_avg, metrics_avg, pred.reshape(*y_data.shape)


# Build model
sep('Model')
model = MLPsb(x_test.shape[1], [(128, 'sigmoid'), (64, 'sigmoid'), (32, 'sigmoid'), (10, 'linear')], LR)

# Train or load
curve_save_path = None
if CKPT_PATH is not None:
    sep('Load weight')
    if not os.path.exists(CKPT_PATH):
        raise Exception(f'Check point file {CKPT_PATH} is not existed')

    print(f'Loading from {CKPT_PATH} ...')
    model.load_weights(CKPT_PATH)
    curve_save_path = get_curve_path(CKPT_PATH)
    print('Loaded.')

if IS_TRAIN:
    sep('Train')
    cost_his, metrics_his = [], []
    cost_his_val, metrics_his_val = [], []
    for i in range(N_EPOCH):
        epoch = i + 1
        cost_avg, metrics_avg, *_ = run_on_data(model, 'train', x_train, y_train_oh, N_BATCH_SIZE, True, epoch)
        metrics_avg = metrics_avg.tolist()
        cost_his.append(cost_avg)
        metrics_his.append(metrics_avg)
        cost_avg_val, metrics_avg_val, *_ = run_on_data(model, 'val', x_val, y_val_oh, N_BATCH_SIZE, False, epoch)
        metrics_avg_val = metrics_avg_val.tolist()
        cost_his_val.append(cost_avg_val)
        metrics_his_val.append(metrics_avg_val)
        save_path = None
        if epoch % SAVE_FREQ == 0:
            save_prefix = get_saved_path_prefix(epoch)
            save_path = get_saved_path(epoch)
            print(f'Saving to {save_path} ...')
            model.save_weights(save_prefix)
            print('Saved.')
        ts.save_epoch(epoch, {
            'cost': cost_avg,
            'acc': metrics_avg[0][0],
            'cost_val': cost_avg_val,
            'acc_val': metrics_avg_val[0][0],
        }, save_path)
    metrics_his = np.array(metrics_his)
    metrics_his_val = np.array(metrics_his_val)

    # save training curve
    plt.figure(figsize=[12, 6])
    spr = 1
    spc = 2
    spn = 0

    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title(f'Cost(lr={LR})')
    plt.plot(cost_his, label='train')
    plt.plot(cost_his_val, label='val')
    plt.legend()
    plt.grid()

    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title('Metrics')
    plt.plot(metrics_his[:, 0], label='acc_train')
    plt.plot(metrics_his_val[:, 0], label='acc_val')
    plt.legend()
    plt.grid()

    save_path = get_saved_path(epoch)
    curve_save_path = get_curve_path(save_path)
    plt.savefig(curve_save_path)
    plt.close()

# do test
sep('Test')
cost_avg, metrics_avg, pred = run_on_data(model, 'test', x_test, y_test_oh, N_BATCH_SIZE, False, 0)

# Draw testing demo in memory buf
plt.figure(figsize=[6, 6])
spr = 5
spc = 5
spn = 0
m = spr * spc
for i in range(m):
    spn += 1
    plt.subplot(spr, spc, spn)
    y = y_test[i, 0]
    img = x_test[i].reshape(*shape_)
    pr = pred[i].argmax()
    plt.axis('off')
    plt.title(f'{y}=>{pr}', color='black' if pr == y else 'red')
    plt.imshow(img)
buf = io.BytesIO()
buf.seek(0)
plt.savefig(buf, format='png')
img_demo = Image.open(buf)
plt.close()

# Show curve of cost and metrics in training if it exits and show the demo of test
if not IS_TRAIN and curve_save_path and os.path.exists(curve_save_path):
    plt.figure(figsize=[15, 5])
    plt.axis('off')
    img2 = Image.open(curve_save_path)
    images = (img_demo, img2)
    widths, heights = zip(*(i.size for i in images))
    total_width = sum(widths)
    max_height = max(heights)
    new_im = Image.new('RGB', (total_width, max_height))
    x_offset = 0
    for im in images:
        new_im.paste(im, (x_offset, 0))
        x_offset += im.size[0]
    plt.imshow(new_im)
    plt.show()
elif IS_TRAIN:
    plt.figure(figsize=[15, 5])
    plt.axis('off')
    spr = 1
    spc = 3
    spn = 0

    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title(f'Cost(lr={LR})')
    plt.plot(cost_his, label='train')
    plt.plot(cost_his_val, label='val')
    plt.legend()
    plt.grid()

    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title('Metrics')
    plt.plot(metrics_his[:, 0], label='acc_train')
    plt.plot(metrics_his_val[:, 0], label='acc_val')
    plt.legend()
    plt.grid()

    spn += 1
    plt.subplot(spr, spc, spn)
    plt.axis('off')
    plt.imshow(img_demo)
    buf.close()

    plt.show()
else:
    plt.figure(figsize=[6, 6])
    plt.axis('off')
    plt.imshow(img_demo)
    buf.close()
    plt.show()
